# Copyright (c) 2021 Presto Labs Pte. Ltd.
# Author: xguo

import json
import os
import pathlib

from cc.appcoin2.strategy.sim_backtest import SimBacktest
from cc.appcoin2.strategy.umm.support.util import (
    get_prefix,
    StratConfig,
)


class UmmBacktest(SimBacktest):
  def __init__(
      self,
      tmp_dirname,
      analyzer_feed_machine,
      *,
      walk_step_bp=None,
      ohlc_window_size_sec=None,
      ohlc_max_len=None,
      rate_window_size_sec=None,
      max_nofill_window_size_sec=None,
      risk_threshold=None,
      take_profit_bp=None,
      stop_loss_bp=None,
      force_fire_cooldown_sec=None,
      use_agg=None,
      price_pull_bp=None,
      price_push_bp=None,
      stack=None,
      sticky_bp=None,
      thrift_ignore_bp=None,
      thrift_ignore_levels=None,
      pick_symbol_regex=None,
      worker=None,
      machine=None,
      **kwargs):
    super().__init__(tmp_dirname, analyzer_feed_machine, **kwargs)

    self.params = {
        "walk_step_bp": walk_step_bp,
        "ohlc_window_size_sec": ohlc_window_size_sec,
        "ohlc_max_len": ohlc_max_len,
        "rate_window_size_sec": rate_window_size_sec,
        "max_nofill_window_size_sec": max_nofill_window_size_sec,
        "risk_threshold": risk_threshold,
        "take_profit_bp": take_profit_bp,
        "stop_loss_bp": stop_loss_bp,
        "force_fire_cooldown_sec": force_fire_cooldown_sec,
        "use_agg": use_agg,
        "price_pull_bp": price_pull_bp,
        "price_push_bp": price_push_bp,
        "stack": stack,
        "sticky_bp": sticky_bp,
        "thrift_ignore_bp": thrift_ignore_bp,
        "thrift_ignore_levels": thrift_ignore_levels,
    }
    self.params = {k: self.params[k] for k in sorted(self.params)}
    self.pick_symbol_regex = pick_symbol_regex
    self.worker = worker or 0
    self.params.update({'config_file': self.config_path})
    self.strat_prefix = get_prefix(self.params)
    self.sim_binary = 'bazel-bin/cc/appcoin2/strategy/umm/umm_sim_runner'

    self.driver_config_filenames = {}
    self.strat_config_filename = {}
    self.umm_strat_config = None

  def gen_strat_config(self, trading_interval, i):
    os.makedirs(f'{self.sim_result_dir}/configs', exist_ok=True)
    st = trading_interval[0].strftime('%Y%m%dT%H%M%S')
    et = trading_interval[1].strftime('%Y%m%dT%H%M%S')
    self.umm_strat_config = StratConfig(
        config_path=self.config_path, is_sim=True, interval=[st, et], sim_dir=self.tmp_dirname)
    self.umm_strat_config.update_params(self.params)
    self.umm_strat_config.gen_config()

    new_config_filename = f'{self.sim_result_dir}/configs/config_{self.strat_prefix}.json'
    with open(new_config_filename, 'w') as fo:
      json.dump(self.umm_strat_config.config, fo, indent=2)

    self.driver_config_filenames[i] = str(self.umm_strat_config.app_config_path)
    self.strat_config_filename[i] = new_config_filename

  def cmd_prerun_on_date(self, trading_interval, i):
    self.gen_strat_config(trading_interval, i)

  def cmd_run_on_date(self, trading_interval, i):
    if self.config_path is None:
      return

    from_dt_str = trading_interval[0].strftime("%Y%m%dT%H%M%S")
    cmd_sim = f"""{self.sim_binary} \
--driver={self.driver_config_filenames[i]} \
--config={self.strat_config_filename[i]} \
&> {self.tmp_dirname}/log{from_dt_str}.txt"""
    print(cmd_sim)
    cmdpath = pathlib.Path(f'{self.sim_result_dir}/log_runsim_cmd.txt')
    if cmdpath.is_file():
      content = cmdpath.read_text() + cmd_sim + '\n'
    else:
      content = cmd_sim + '\n'
    cmdpath.write_text(content)
    return cmd_sim
